#include "Containers/RangePair.h"
#include "File/BOSSOutputLoader.h"
#include "File/YamlLoader.h"
#include "Functions/Draw.h"
#include "Functions/File.h"
#include "Functions/PrintTable.h"
#include "Globals/Settings.h"
#include "TArrow.h"
#include "TCanvas.h"
#include "TCut.h"
#include "TLatex.h"
#include "TLegend.h"
#include "TStyle.h"
#include <iostream>

using namespace BOSS_Afterburner;
using namespace BOSS_Afterburner::Draw;
using namespace BOSS_Afterburner::Error;
using namespace BOSS_Afterburner::File;
using namespace RooFit;
using namespace BOSS_Afterburner::Detector::Units;
using namespace std;

void Chi2Optimisation(const char* config_file)
{
  YamlLoader conf(config_file);

  // * Data files * //
  const auto& files = conf["Input files"];
  if(!files) return;
  BOSSOutputLoader excl(files["excl"].as<string>().c_str());
  BOSSOutputLoader incl(files["incl"].as<string>().c_str());
  BOSSOutputLoader data(files["data"].as<string>().c_str());

  // * Tree name * //
  const auto& treeDef = conf["Tree name"];
  if(!treeDef) return;
  string tree = treeDef.as<string>().c_str();

  // * Determine cuts * //
  TCut       massCut;
  const auto cutMap = conf["Mass cut"];
  if(cutMap.IsSequence())
    for(const auto& cutDef : cutMap)
      massCut += cutDef.as<string>().c_str();

  // * Determine tree name * //
  const auto& plotDef = conf["Chi2 plots"];
  if(!plotDef) return;

  // * Get chi squared plots * //
  auto exclHist = excl.GetChain(tree).DrawHistogram(plotDef["excl"]);
  auto inclHist = incl.GetChain(tree).DrawHistogram(plotDef["incl"]);
  auto dataHist = data.GetChain(tree).DrawHistogram(plotDef["data"]);

  // * Normalise histograms * //
  if(conf["Normalise"].IsDefined() && conf["Normalise"].as<bool>())
  {
    exclHist->Scale(1. / exclHist->GetBinContent(exclHist->GetMaximumBin()));
    inclHist->Scale(1. / inclHist->GetBinContent(inclHist->GetMaximumBin()));
    dataHist->Scale(1. / dataHist->GetBinContent(dataHist->GetMaximumBin()));
  }

  // * Beautify histograms * //
  inclHist->Sumw2(false);
  inclHist->SetLineWidth(0.);
  inclHist->SetFillStyle(1001);
  inclHist->SetFillColor(kGray);
  inclHist->SetFillColorAlpha(kGray + 2, .25);
  exclHist->SetLineColor(kRed);

  // * Get ranges of plots * //
  double  x1, x2, y1, y2, max = 0.;
  TCanvas temp;
  inclHist->Draw("e1");
  temp.Update();
  temp.GetRangeAxis(x1, y1, x2, y2);
  if(y2 > max) max = y2;
  exclHist->Draw("e1");
  temp.Update();
  temp.GetRangeAxis(x1, y1, x2, y2);
  if(y2 > max) max = y2;
  dataHist->Draw("e1");
  temp.Update();
  temp.GetRangeAxis(x1, y1, x2, y2);
  if(y2 > max) max = y2;

  // * Draw chi squared plots * //
  TCanvas c;
  exclHist->GetYaxis()->SetRangeUser(0., max);
  inclHist->GetYaxis()->SetTitle("#it{N} (a.u.)");
  inclHist->GetYaxis()->SetLabelSize(0.);
  inclHist->GetYaxis()->SetNdivisions(1);
  inclHist->Draw();
  dataHist->Draw("same");
  exclHist->Draw("e1 same");
  inclHist->SetTitle("incl");
  dataHist->SetTitle("data");
  exclHist->SetTitle("excl");
  gStyle->SetLegendTextSize(.04);
  c.BuildLegend(.85, .14, .94, .38);
  c.Update();
  SaveCanvas("chi2", &c);

  // * Perform optimal chi2 cut scan * //
  const auto& scanDef = conf["Chi2 scan"];
  if(!scanDef) return;
  double      chiMax  = 200.;
  double      step    = 1.;
  const auto& maxDef  = scanDef["max"];
  const auto& stepDef = scanDef["step size"];
  if(maxDef.IsScalar()) chiMax = maxDef.as<double>();
  if(stepDef.IsScalar()) step = stepDef.as<double>();

  auto Nexcl = exclHist->GetEntries();
  auto Nincl = inclHist->GetEntries();

  Int_t  nbins = floor(chiMax / step) + 1;
  TGraph exclEffGraph(nbins);
  TGraph inclEffGraph(nbins);
  TGraph yieldGraph(nbins);

  exclEffGraph.SetLineColor(kRed);
  inclEffGraph.SetLineColor(kBlue);
  yieldGraph.SetLineColor(kGreen + 2);

  const char* epsSig = "#epsilon_{sig}";
  const char* epsBck = "#epsilon_{bck}";
  exclEffGraph.SetTitle(epsSig);
  inclEffGraph.SetTitle(epsBck);
  exclEffGraph.GetXaxis()->SetTitle(exclHist->GetXaxis()->GetTitle());
  yieldGraph.SetTitle(Form("#frac{%s}{#sqrt{%s + %s}}", epsSig, epsSig, epsBck));

  Table    tbl;
  Double_t bestChi2  = 0.;
  Double_t bestYield = 0.;
  for(Double_t chi = 0.; chi <= chiMax; chi += step)
  {
    Int_t   i       = round(chi / step);
    string chiName = "chisq";
    try
    {
      chiName = conf["Chi name"].as<string>();
    }
    catch(const YAML::BadConversion& e)
    {}
    TCut cut = massCut + Form("%s < %g", chiName.c_str(), chi);
    excl.GetChain(tree).Draw(Form("%s>>excl", chiName.c_str()), cut);
    auto exclChiHist = dynamic_cast<TH1F*>(gPad->GetPrimitive("excl"));
    incl.GetChain(tree).Draw(Form("%s>>incl", chiName.c_str()), cut);
    auto inclChiHist = dynamic_cast<TH1F*>(gPad->GetPrimitive("incl"));
    if(!exclChiHist && !inclChiHist) continue;
    Double_t effSig = exclChiHist->GetEntries() / (Double_t)Nexcl;
    Double_t effBck = inclChiHist->GetEntries() / (Double_t)Nincl;
    Double_t yield  = effSig / sqrt(effSig + effBck);
    if(yield > bestYield)
    {
      bestYield = yield;
      bestChi2  = chi;
    }
    if(yield != yield) yield = 0.;
    exclEffGraph.SetPoint(i, chi, effSig);
    inclEffGraph.SetPoint(i, chi, effBck);
    yieldGraph.SetPoint(i, chi, yield);
    tbl << chi << effSig << effBck << yield << TableRow::endrow;
    delete exclChiHist, inclChiHist;
  }

  tbl.SetHeader({"chi2", "effSig", "effBck", "yield"});
  tbl.SetAlign({TableColumn::Right});
  tbl.Print();
  tbl.WriteLaTeX("yields");
  tbl.Write("yields");

  // * Draw optimal chi2 cut scan * //
  TCanvas c2;
  exclEffGraph.Draw();
  inclEffGraph.Draw("same");
  yieldGraph.Draw("same");
  exclEffGraph.GetXaxis()->SetRangeUser(0., chiMax);
  Draw::FormatLegend(c2.BuildLegend(), scanDef);

  // * Indicate best chi2 * //
  TArrow arrow(bestChi2, bestYield + .1, bestChi2, bestYield);
  TLatex text(bestChi2, bestYield + .12, Form("#chi^{2}_{best} = %g", bestChi2));
  text.SetTextColor(kGreen + 2);
  text.SetTextAlign(21);
  arrow.SetArrowSize(.02);
  arrow.SetLineColor(kGreen + 2);
  arrow.Draw();
  text.Draw();

  // * Draw optimal chi2 cut scan * //
  Draw::SaveCanvas("yields", &c2);
}

/// Main function that is called when executing the executable compiled using e.g. `g++`.
int main(int argc, char* argv[])
{
  TerminalIO::CheckMainArguments(argc, argv, 1);
  Chi2Optimisation(argv[1]);
  return 0;
}